
#include <iostream>
#include <fstream>
#include <stdio.h>
#include "frontend/Feature.h"
#include <cmath>

// #define MNN_INFER

#ifdef ONNXRUNTIME_INFER
    #include "kws/keyword_spotting.h"
    using namespace wekws;
#elif defined(MNN_INFER)
    #include "kws/keywordspotter.h"
    using namespace tykws;
#else
    #error "must choose ONNXRUNTIME_INFER or MNN_INFER"
#endif

// 特征提取测试
#define T 300  // ms
#define CHUNK (T*16*2)

int main(int argc, char *argv[])
{
    std::shared_ptr<wenet::FeaturePipeline> g_feature_pipeline;
    wenet::MODEL_TYPE mode_type;
    mode_type = wenet::CTC_TYPE_MODEL;
    int num_bins = 80;
    wenet::FeaturePipelineConfig feature_config(num_bins, 16000, mode_type);
    g_feature_pipeline = std::make_shared<wenet::FeaturePipeline>(feature_config);

    std::string token_path;
    std::string key_word;
    const std::string model_path = argv[1];
    token_path = argv[2];     
    key_word = "空气很好";
    KeywordSpotting spotter(model_path, DECODE_PREFIX_BEAM_SEARCH, mode_type);
    spotter.readToken(token_path);
    spotter.setKeyWord(key_word);


    std::ifstream fidr(argv[3], std::ios::binary); // read pcm data
    // 获取文件大小
    fidr.seekg(0, std::ifstream::end);
    size_t fileSize = fidr.tellg();
    fidr.seekg(0, std::ifstream::beg);
    int nchunks = fileSize/CHUNK;
    std::cout<<"total chunks = "<<nchunks<<std::endl;
    char wavbuffer[CHUNK] = {0};

    FILE *fidw = fopen("prob_mnn.bin", "wb");

    for (int n=0; n<nchunks; n++) 
    {
        // std::cout<<"frame num = "<<n<<std::endl;
        fidr.read(wavbuffer, CHUNK);
        int16_t *pcm_data = (int16_t *)wavbuffer;
        std::vector<int16_t> v(pcm_data, pcm_data+16*T);
        g_feature_pipeline->AcceptWaveform(v);        
        std::vector<std::vector<float>> feats;
        g_feature_pipeline->Read(9, &feats);
        // printf("feats = [%ld ✖ %ld]\n", feats.size(), feats[0].size());         
        
        if (n>=1) {
            // for (int i=0; i<feats.size(); i++) {
            //     for (int j=0; j<feats[0].size(); j++) {
            //         fwrite(&feats[i][j], sizeof(float), 1, fidw);
            //     }
            // }
            std::vector<std::vector<float>> probs;
            spotter.Forward(feats, &probs);
            // printf("probs = [%ld ✖ %ld]\n", probs.size(), probs[0].size());
            for (int i=0; i<probs.size(); i++) {
                for (int j=0; j<probs[0].size(); j++) {
                    fwrite(&probs[i][j], sizeof(float), 1, fidw);
                }
            }
            
            if (mode_type==wenet::CTC_TYPE_MODEL) {
                float hitScoreThr = 0.1f;    // threshold of hit score.
                spotter.decode_keywords(probs, hitScoreThr);
            }
        }        
    }

    fclose(fidw);
    fidr.close();
    std::cout<<"keyword spot testing...."<<std::endl;
    return 0;
}

